# VideoThumbnailer - Generates thumbnails from video files
# Copyright (C) 2007 Paul Hummer
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

__license__ = "GPLv2"
__copyright__ = "2007, Paul Hummer"
__author__ = "Paul Hummer <paul@ironlionsoftware.no.spam.com>"

# I think it's important to note that the methodology of this code (and 
# some actual EXACT snippets of code) were based on Elisa (http://elisa.fluendo.com/)

import os
import sys
from threading import Event
try:
    import hashlib
except ImportError:
    import md5

try:
    import gobject
    import pygtk
    import gtk
except ImportError:
    print "The VideoThumbnailer requires the gtk python library"
    sys.exit(1)

try:
    import pygst
    pygst.require('0.10')
    import gst
except ImportError:
    print "The VideoThumbnailer requires the gtk gstreamer python library"
    sys.exit(1)

try:
    import Image
    import ImageStat
except ImportError:
    print "The VideoThumbnailer requires the python imaging library"
    sys.exit(1)

import thumbnailer

class VideoThumbnailerException(Exception): pass

class VideoThumbnailer(thumbnailer.Thumbnailer):
    ''' VideoThumbnailer does exactly what it's name demonstrates it does:
    create thumbnailed images of the videos.'''

    class VideoSinkBin(gst.Bin):
    
        def __init__(self, needed_caps):
            self.reset()
            gst.Bin.__init__(self)
            self._capsfilter = gst.element_factory_make('capsfilter', 'capsfilter')
    
            self.set_caps(needed_caps)
            self.add(self._capsfilter)
    
            fakesink = gst.element_factory_make('fakesink', 'fakesink')
            fakesink.set_property("sync", False)
            self.add(fakesink)
            self._capsfilter.link(fakesink)
    
            pad = self._capsfilter.get_pad("sink")
            ghostpad = gst.GhostPad("sink", pad)
    
            pad2probe = fakesink.get_pad("sink")
            pad2probe.add_buffer_probe(self.buffer_probe)
    
            self.add_pad(ghostpad)
            self.sink = self._capsfilter
    
        def set_current_frame(self, value):
            self._current_frame = value
    
        def set_caps(self, caps):
            gst_caps = gst.caps_from_string(caps)
            self._capsfilter.set_property("caps", gst_caps)
    
        def get_current_frame(self):
            frame = self._current_frame
            self._current_frame = None
            return frame
    
        def buffer_probe(self, pad, buffer):
            caps = buffer.caps
            if caps != None:
                s = caps[0]
                self.width = s['width']
                self.height = s['height']
            if self.width != None and self.height != None and buffer != None:
                self.set_current_frame(buffer.data)
            return True
    
        def reset(self):
            self.width = None
            self.height = None
            self.set_current_frame(None)
            

    def __init__(self, filename, src="VIDEO"):

        if os.path.exists(filename):
            if os.path.isfile(filename):
                self._filename = filename
                self._fileuri = 'file://%s' % (self._filename)
            else:
                raise VideoThumbnailerException('VideoThumbnailer requires a file to be passed as second argument.  A folder was passed instead')
        else:
            raise VideoThumbnailerException('VideoThumbnailer filename parameter does not exist')

        if src == "VIDEO":
            dest_dir = os.path.expanduser("~/.entertainer/cache/thumbnails/video")
        elif src == "RECORDING":
            dest_dir = os.path.expanduser("~/.entertainer/cache/thumbnails/recording")
        else:
            raise VideoThumbnailerException('VideoThumbnailer was provided an invalid source type.')

        self.THUMB_SIZE = 500
        self.BORING_THRESHOLD = 2000
        self.HOLES_SIZE = (9, 35)
        self.HOLES_DATA = '\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6'
        
        if hashlib:
            m = hashlib.md5()
        else:
            m = md5.new()
        m.update(self._filename)
        thumb_file = '%s_%s.jpg' % (m.hexdigest(), self.THUMB_SIZE)
        self._thumb_file = os.path.join(dest_dir, thumb_file)
        
        #Initialize and use the gstreamer pipeline
        self._pipeline = gst.element_factory_make('playbin', 'playbin')
        self._sink = self.VideoSinkBin('video/x-raw-rgb,bpp=24,depth=24')
        self._blocker = Event()
        self._pipeline.set_property('video-sink', self._sink)
        self._pipeline.set_property('volume', 0)

    def add_holes(self, img):
        holes = Image.fromstring('RGBA', self.HOLES_SIZE, self.HOLES_DATA)
        holes_h = holes.size[1]
        remain = img.size[1] % holes_h

        i = 0
        nbands = 0
        while i < (img.size[1] - remain):
            left_box = (0, i, holes.size[0], (nbands+1) * holes.size[1])
            img.paste(holes, left_box)

            right_box = (img.size[0] - holes.size[0], i,
                         img.size[0], (nbands+1) * holes.size[1])
            img.paste(holes, right_box)

            i += holes_h
            nbands += 1

        remain_holes = holes.crop((0, 0, holes.size[0], remain))
        remain_holes.load()
        img.paste(remain_holes, (0, i, holes.size[0], img.size[1]))
        img.paste(remain_holes, (img.size[0] - holes.size[0], i,
                                 img.size[0], img.size[1]))
        return img
    
    def interesting_image(self, img):
        '''
        Checks an image to see if it has the characteristics of an 
        'intersting' image, i.e. whether or not the image is worth 
        examining for thumbnailing
        '''
        stat = ImageStat.Stat(img)
        return True in [ i > self.BORING_THRESHOLD for i in stat.var ]

    def set_pipeline_state(self, pipeline, state):
        '''
        Does exactly what it claims : sets the state of the pipeline, and returns
        the boolean value of whether or not is successfully changed state
        '''
        status = pipeline.set_state(state)
        if status == gst.STATE_CHANGE_ASYNC:

            result = [False]
            max_try = 100
            nb_try = 0
            while not result[0] == gst.STATE_CHANGE_SUCCESS:
                if nb_try > max_try:
                    #State change failed
                    return False
                nb_try += 1
                result = pipeline.get_state(50*gst.MSECOND)

            return True
        elif status == gst.STATE_CHANGE_SUCCESS:
            return True
        else:
            return False


    def create_thumbnail(self):
        '''
        Creates the new thumbnail based on the information provided through the
        gstreamer API.  The file is an md5 of the video's file name, followed 
        by the .jpg extension
        '''
        self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
        self._pipeline.set_property('uri', self._fileuri)

        if not self.set_pipeline_state(self._pipeline, gst.STATE_PAUSED):
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            raise VideoThumbnailerException('Cannot start the pipeline')

        if self._sink.width == None or self._sink.height == None:
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            raise VideoThumbnailerException('Unable to determine media size')
        sink_size = (self._sink.width, self._sink.height)

        try:
            duration, format = self._pipeline.query_duration(gst.FORMAT_TIME)
        except Exception, e:
        ## FIXME: precise this exception
            #Gstreamer cannot determine the media duration using playing-thumbnailing for file
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            img = self._play_for_thumb(sink_size, size, 0)
            #play found img
            if img:
                img.save(self._thumb_file)
                return
        else:
            duration /= gst.NSECOND
            try:
               img = self._seek_for_thumb(duration, sink_size) 
               #Seek found image
               if img:
                   img.save(self._thumb_file)
                   return
            except VideoThumbnailerException, e:
                # Fallback: No Image found in seek_for, falling back to play_for_thumb
                self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
                img = self._play_for_thumb(sink_size, duration)
                #Fallback-Play found img
                if img:
                    img.save(self._thumb_file)
                    return

        self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
        raise VideoThumbnailerException('Unable to create thumbnail.  Please file a bug')


    def _play_for_thumb(self, sink_size, duration=0):
        '''
        Plays the video file to gather information for generating a thumbnail
        '''
        id = None
        self._img = None

        if duration >= 250000:
            self._every = 25
        elif duration >= 200000:
            self._every = 15
        elif duration >= 10000:
            self._every = 10
        elif duration >= 5000:
            self._every = 5
        else:
            self._every = 1

        #Setting every-frame to self._every

        self._every_co = self._every

        ## How often Proceed?
        self._counter = 5
        
        def buffer_probe(pad, buffer):
            ## Proceed only every 5th frame!
            if self._every_co < self._every:
                self._every_co += 1
                return
            self._every_co = 0

            try: 
                img = Image.frombuffer("RGB", sink_size, buffer,
                                                "raw", "RGB",0, 1)
            except Exception, e:
                #Invalid frame
                pass
            else:
                self._img = img
                if self.interesting_image(self._img):

                    self._img.thumbnail((self.THUMB_SIZE, self.THUMB_SIZE), Image.BILINEAR)
                    if img.mode != 'RGBA':
                        img = img.convert(mode='RGBA')

                    #Releasing self._img
                    self._sink.reset()
                    pad.remove_buffer_probe(id)
                    self._blocker.set()
                    return

            self._counter -= 1
            if self._counter <= 0:
                #Counter off, resetting blocker
                # Is it better to return no image instead of a 'boring' one?
                if self._img:
                    self._img.thumbnail((size, size), Image.BILINEAR)
                    if img.mode != 'RGBA':
                        img = img.convert(mode='RGBA')

                self._sink.reset()
                pad.remove_buffer_probe(id)
                self._blocker.set()


        self.set_state_blocking(self._pipeline, gst.STATE_PLAYING)

        pad = self._sink.get_pad('sink')
        id = pad.add_buffer_probe(buffer_probe)
        self._blocker.wait()
        self._pipeline.set_state(gst.STATE_NULL)
        return self._img

    def _seek_for_thumb(self, duration, sink_size):
        '''
        Seeks through the video file to gather information for generating a thumbnail 
        '''
        frame_locations = [ 1.0 / 3.0, 2.0 / 3.0, 0.1, 0.9, 0.5 ]

        for location in frame_locations:
            abs_location = int(location * duration)

            if abs_location == 0:
                raise VideoThumbnailerException(self._filename, 'Video has a size of zero')
            
            event = self._pipeline.seek(1.0, gst.FORMAT_TIME,
                                        gst.SEEK_FLAG_FLUSH | gst.SEEK_FLAG_KEY_UNIT,
                                        gst.SEEK_TYPE_SET, abs_location,
                                        gst.SEEK_TYPE_NONE, 0)
            if not event:
                raise VideoThumbnailerException(self._filename, 'Unable to seek through video')

            if not self.set_pipeline_state(self._pipeline, gst.STATE_PAUSED):
                raise VideoThumbnailerException(self._filename, 'Unable to pause video')

            frame = self._sink.get_current_frame()

            try:
                img = Image.frombuffer("RGB", sink_size, frame, "raw", "RGB", 0, 1)
            except:
                #Invalid frame
                continue

            if self.interesting_image(img):
                break
            else:
                pass

        self._sink.reset()

        if img:
            img.thumbnail((self.THUMB_SIZE, self.THUMB_SIZE), Image.BILINEAR)
            if img.mode != 'RGBA':
                img = img.convert(mode='RGBA')
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            return img

gobject.type_register(VideoThumbnailer.VideoSinkBin)
